from docplex.mp.model import Model
from QKDNetwork import QKDNetwork

def getLoadBalanceRouteList(net: QKDNetwork) -> list: # 使用负载均衡求解候选路径 混合中继 Version
    # return [{"sd": [1, 2, 4], "path": [1, 3, 2]}...]
    model = Model(name="load_balance")
    alpha = model.continuous_var(lb=0, name="alpha")
    f_dim1 = net.num_nodes + 1
    f = model.continuous_var_dict([i for i in range(f_dim1)], lb=0, name="f")
    q_dim1, q_dim2 = len(net.sd_list), net.max_candidate_size
    q = model.binary_var_dict([(i, j) for i in range(q_dim1) for j in range(q_dim2)], lb=0, name="q" )
    
    # 6c ✅ 所有节点的负载小于 系数 * 速率
    for u in net.G.nodes:
        if(net.G.nodes[u]["transmitter"] == 1):
            model.add_constraint(f[u] <= alpha * net.G.nodes[u]["transmitter"])
    
    # 6d ✅ 
    for u in net.G.nodes:
        if(net.G.nodes[u]["transmitter"] == 1):
            flow_expr = model.linear_expr()
            flag = False
            for sd_index, sd_item in enumerate(net.sd_list):
                for candidate_path_index, candidate_path_item in enumerate(net.candidate_paths_mdi[sd_index]["paths"]):
                    if(u in candidate_path_item):
                        flag = True
                        flow_expr += q[sd_index, candidate_path_index] 
            if(flag):
                flow_expr -= f[u]
                model.add_constraint(flow_expr <= 0)
    
    # 6c ✅
    for sd_index, sd_item in enumerate(net.sd_list):
        at_least_one_path = model.linear_expr()
        at_least_one_path += 1
        for candidate_path_index, candidate_path_item in enumerate(net.candidate_paths_mdi[sd_index]["paths"]):
            at_least_one_path -= q[sd_index, candidate_path_index]
        model.add_constraint(at_least_one_path == 0)
    
    model.minimize(alpha)
    sol = model.solve()
    ret = []
    if sol:
        # print(model.get_var_by_name("alpha").solution_value)
        for (i, j), variable in q.items():
            if(variable.solution_value == 1):
                ret.append({
                    "sd": net.sd_list[i],
                    "path": net.candidate_paths_mdi[i]["paths"][j]
                })
                # print(f"q[{i}, {j}] = {variable.solution_value}")
    return ret
